import json
from typing import Any, Union, List, Optional
from functools import wraps
from utils.logger import logger


class AssertionError(Exception):
    """自定义断言异常"""
    def __init__(self, path: str, expected: Any, actual: Any, message: str):
        self.path = path
        self.expected = expected
        self.actual = actual
        self.message = message
        super().__init__(self._format_message())

    def _format_message(self) -> str:
        return (f"Assertion Failed at '{self.path}':\\n"
                f"Expected: {self._format_value(self.expected)}\\n"
                f"Actual  : {self._format_value(self.actual)}\\n"
                f"Message : {self.message}")

    @staticmethod
    def _format_value(value: Any) -> str:
        if isinstance(value, (dict, list)):
            return json.dumps(value, indent=2, ensure_ascii=False)
        return str(value)

class ResponseAssertion:
    """链式断言处理器"""

    def __init__(self,
                 target: Any,
                 path: str = 'root',
                 parent: Optional['ResponseAssertion'] = None):
        self._target = target
        self._path = path
        self._parent = parent

    def __getattr__(self, name: str) -> 'ResponseAssertion':
        """处理对象属性访问"""
        return self.key(name)

    def __getitem__(self, index: int) -> 'ResponseAssertion':
        """处理数组索引访问"""
        return self.index(index)

    def key(self, key: str) -> 'ResponseAssertion':
        """访问字典键"""
        try:
            new_target = self._target[key]
            new_path = f"{self._path}.{key}"
        except KeyError:
            self._raise_missing(f"Key '{key}' not found")
        except TypeError:
            self._raise_type_error("dict", type(self._target))
        return ResponseAssertion(new_target, new_path, self)

    def index(self, index: int) -> 'ResponseAssertion':
        """访问数组索引"""
        try:
            new_target = self._target[index]
            new_path = f"{self._path}[{index}]"
        except IndexError:
            self._raise_missing(f"Index {index} out of range")
        except TypeError:
            self._raise_type_error("list", type(self._target))
        return ResponseAssertion(new_target, new_path, self)

    def length(self) -> 'ResponseAssertion':
        """获取数组/字典长度"""
        try:
            return ResponseAssertion(len(self._target), f"{self._path}.length", self)
        except TypeError:
            self._raise_type_error("list/dict", type(self._target))

    def _raise_missing(self, msg: str):
        """抛出字段缺失异常"""
        raise AssertionError(
            path=self._path,
            expected="Existing field",
            actual="Missing field",
            message=msg
        )

    def _raise_type_error(self, expected_type: str, actual_type: type):
        """抛出类型错误异常"""
        raise AssertionError(
            path=self._path,
            expected=f"{expected_type} type",
            actual=f"{actual_type.__name__} type",
            message="Type mismatch"
        )

    def should_equal(self, expected: Any) -> 'ResponseAssertion':
        """值相等断言"""
        if self._target != expected:
            raise AssertionError(
                path=self._path,
                expected=expected,
                actual=self._target,
                message="Values are not equal"
            )
        return self

    def should_contain(self, item: Any) -> 'ResponseAssertion':
        """包含关系断言"""
        if item not in self._target:
            raise AssertionError(
                path=self._path,
                expected=f"Contains {item}",
                actual=self._target,
                message="Item not found"
            )
        return self

    # def should_match_schema(self, schema: dict) -> 'ResponseAssertion':
    #     """JSON Schema验证"""
    #     from jsonschema import validate, ValidationError
    #     try:
    #         validate(instance=self._target, schema=schema)
    #     except ValidationError as e:
    #         raise AssertionError(
    #             path=self._path,
    #             expected=schema,
    #             actual=self._target,
    #             message=f"Schema validation failed: {e.message}"
    #         )
    #     return self

    def log_on_error(self):
        """错误日志装饰器"""
        def decorator(func):
            @wraps(func)
            def wrapper(*args, **kwargs):
                try:
                    return func(*args, **kwargs)
                except AssertionError as e:
                    logger.error(e)
                    raise
            return wrapper
        return decorator

# 使用示例
if __name__ == "__main__":
    response = {
        "code": 200,
        "data": {
            "users": [
                {"id": 1, "name": "Alice"},
                {"id": 2, "name": "Bob"}
            ],
            "pagination": {
                "page": 1,
                "page_size": 20
            }
        }
    }

    # 链式断言示例
    try:
        (ResponseAssertion(response)
         .key('code').should_equal(200)
         .key('data').key('users').length().should_equal(2)
         .key('data').key('users')[0].key('id').should_equal(1)
         .key('data').key('pagination').key('page_size').should_equal(20))
    except AssertionError as e:
        logger.error(e)
        raise
